package com.joe.lucene.facet.salmonrun;

import java.util.BitSet;
import java.util.HashMap;
import java.util.Map;

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.Hits;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.QueryFilter;

public class BitSetFacetHitCounter {

	private Query baseQuery;
	private Map<String, Query> subQueries;
	private IndexSearcher searcher;

	public BitSetFacetHitCounter() {
		super();
	}

	public void setBaseQuery(Query baseQuery) {
		this.baseQuery = baseQuery;
	}

	public void setSubQueries(Map<String, Query> subQueries) {
		this.subQueries = subQueries;
	}

	public void setSearcher(IndexSearcher searcher) {
		this.searcher = searcher;
	}

	public Map<String, Integer> getFacetHitCounts() throws Exception {
		Map<String, Integer> facetCounts = new HashMap<String, Integer>();
		IndexReader reader = searcher.getIndexReader();
		QueryFilter baseQueryFilter = new QueryFilter(baseQuery);
		BitSet baseBitSet = baseQueryFilter.bits(reader);
		for (String attribute : subQueries.keySet()) {
			QueryFilter filter = new QueryFilter(subQueries.get(attribute));
			BitSet filterBitSet = filter.bits(reader);
			facetCounts.put(attribute, getFacetHitCount(baseBitSet,
					filterBitSet));
		}
		return facetCounts;
	}

	private int getFacetHitCount(BitSet baseBitSet, BitSet filterBitSet) {
		filterBitSet.and(baseBitSet);
		return filterBitSet.cardinality();
	}

	// using Boolean queries
	public Map<String, Integer> getFacetHitCounts1() throws Exception {
		Map<String, Integer> facetCounts = new HashMap<String, Integer>();
		for (String attribute : subQueries.keySet()) {
			BooleanQuery facetQuery = new BooleanQuery();
			facetQuery.add(baseQuery, BooleanClause.Occur.MUST);
			facetQuery.add(subQueries.get(attribute), BooleanClause.Occur.MUST);
			Hits hits = searcher.search(facetQuery);
			facetCounts.put(attribute, hits.length());
		}
		return facetCounts;
	}

	// using search with QueryFilters
	public Map<String, Integer> getFacetHitCounts2() throws Exception {
		Map<String, Integer> facetCounts = new HashMap<String, Integer>();
		for (String attribute : subQueries.keySet()) {
			QueryFilter filter = new QueryFilter(subQueries.get(attribute));
			Hits hits = searcher.search(baseQuery, filter);
			facetCounts.put(attribute, hits.length());
		}
		return facetCounts;
	}

}
